package com.gt.sql

import com.gt.SCUtil
import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types.{DataType, LongType, StringType, StructField, StructType}
import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}

/**
 * DSL 选择指定的列
 * 首先 -> 引入 import spark.implicits._
 *
 */
object SparkDataFrame_04_UDAF {
  def main(args: Array[String]): Unit = {
    val spark: SparkSession = SCUtil.buildLocalSparkSession()
    import spark.implicits._

    val df: DataFrame = spark.read.json("data/user.json")
    val ds: Dataset[User] = df.as[User]

    spark.udf.register("avgAge",new myUDAF())

    ds.createOrReplaceTempView("user")

    val resultDF: DataFrame = spark.sql("select avgAge(name,age) as avgAge from user")
    resultDF.show()

    spark.close()
  }

  class myUDAF extends UserDefinedAggregateFunction{
    //输入数据结构 in
    override def inputSchema: StructType = StructType(
      Array(
        StructField("name",StringType),
        StructField("age",LongType)
      )
    )

    //缓冲区的数据结构（临时数据的数据结构）buffer
    override def bufferSchema: StructType = StructType(
      Array(
        StructField("total", LongType),
        StructField("count", LongType)
      )
    )

    //数据结果的数据格式 out
    override def dataType: DataType = LongType

    //函数稳定性（函数传入相同数据后的结果是否相同）
    override def deterministic: Boolean = true

    //缓冲区数据初始化
    override def initialize(buffer: MutableAggregationBuffer): Unit = {
      //buffer(0) = 0l
      //buffer(1) = 0l
      buffer.update(0, 0L)
      buffer.update(1, 0L)
    }

    //根据输入数据逐(input)条更新缓冲区数据(buffer)
    override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
      val value: Long = input.getLong(1)
      buffer.update(0, buffer.getLong(0)+value)
      buffer.update(1, buffer.getLong(1)+1)
    }

    //缓冲区数据合并
    override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
      buffer1.update(0, buffer1.getLong(0) + buffer2.getLong(0))
      buffer1.update(1, buffer1.getLong(1) + buffer2.getLong(1))
    }

    //根据缓冲区数据计算结果
    override def evaluate(buffer: Row): Any = {
      buffer.getLong(0) / buffer.getLong(1)
    }
  }


}
